from typing import Optional, Annotated

from fastapi import APIRouter, Query, Depends
from sqlalchemy import select, distinct

from models.models import Forecast, Statistic
from db.db import AsyncSession, get_session
from db.forecast import get_all_forecast_name
from choices.choices import Modality, Locality
from routers.users import oauth2_scheme

forecast_router = APIRouter()


@forecast_router.get("/forecast_name")
async def get_forecast_name(
    token: Annotated[str, Depends(oauth2_scheme)],
    session: AsyncSession = Depends(get_session),
):
    """Возвращает все названия прогноза."""
    return await get_all_forecast_name(session=session)


@forecast_router.get("/hour_forecast_count")
async def get_hour_forecast_count(
    forecast_id: int,
    token: Annotated[str, Depends(oauth2_scheme)],
    modality: Modality = Query(None, description="Модальность"),
    session: AsyncSession = Depends(get_session),
):
    """Возвращает прогноз количества исследований."""
    query = (
        select(Forecast.date)
        .where(
            Forecast.id_filename == forecast_id,
        )
        .distinct(
            Forecast.date,
        )
    )

    if modality:
        query = query.where(Forecast.modality == modality.value)

    query_result = await session.execute(query)
    unique_dates = query_result.scalars().all()

    query = select(Forecast).where(
        Forecast.id_filename == forecast_id,
    )

    if modality:
        query = query.where(Forecast.modality == modality.value)

    query_result = await session.execute(query)
    forecast_datas = query_result.scalars().all()

    modality_dict = {
        "РГ": {},
        "КТ": {},
        "МРТ": {},
        "ММГ": {},
        "Флг": {},
        "Денс": {},
        "РХ": {},
        "НДКТ": {},
        "ПЭТ/КТ": {},
    }

    for el in forecast_datas:
        modality_dict[el.modality][el.date] = (
            modality_dict[el.modality].get(el.date, 0) + el.number
        )
    return {
        "labels": unique_dates,
        "datasets": [
            {
                "label": "РГ",
                "data": [el for el in modality_dict["РГ"].values()],
                "barThickness": 30,
            },
            {
                "label": "КТ",
                "data": [el for el in modality_dict["КТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "МРТ",
                "data": [el for el in modality_dict["МРТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "ММГ",
                "data": [el for el in modality_dict["ММГ"].values()],
                "barThickness": 30,
            },
            {
                "label": "Флг",
                "data": [el for el in modality_dict["Флг"].values()],
                "barThickness": 30,
            },
            {
                "label": "Денс",
                "data": [el for el in modality_dict["Денс"].values()],
                "barThickness": 30,
            },
            {
                "label": "РХ",
                "data": [el for el in modality_dict["РХ"].values()],
                "barThickness": 30,
            },
            {
                "label": "НДКТ",
                "data": [el for el in modality_dict["НДКТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "ПЭТ/КТ",
                "data": [el for el in modality_dict["ПЭТ/КТ"].values()],
                "barThickness": 30,
            },
        ],
    }


@forecast_router.get("/hour_forecast_time")
async def get_hour_forecast_time(
    forecast_id: int,
    token: Annotated[str, Depends(oauth2_scheme)],
    modality: Modality = Query(None, description="Модальность"),
    session: AsyncSession = Depends(get_session),
):
    """Возвращает прогноз времени исследований."""
    query = (
        select(Forecast.date)
        .where(
            Forecast.id_filename == forecast_id,
        )
        .distinct(
            Forecast.date,
        )
    )

    if modality:
        query = query.where(Forecast.modality == modality.value)

    query_result = await session.execute(query)
    unique_dates = query_result.scalars().all()

    query = select(Forecast).where(
        Forecast.id_filename == forecast_id,
    )

    if modality:
        query = query.where(Forecast.modality == modality.value)

    query_result = await session.execute(query)
    forecast_datas = query_result.scalars().all()

    modality_dict = {
        "РГ": {},
        "КТ": {},
        "МРТ": {},
        "ММГ": {},
        "Флг": {},
        "Денс": {},
        "РХ": {},
        "НДКТ": {},
        "ПЭТ/КТ": {},
    }

    for el in forecast_datas:
        modality_dict[el.modality][el.date] = (
            modality_dict[el.modality].get(el.date, 0) + el.workhours
        )
    return {
        "labels": unique_dates,
        "datasets": [
            {
                "label": "РГ",
                "data": [el for el in modality_dict["РГ"].values()],
                "barThickness": 30,
            },
            {
                "label": "КТ",
                "data": [el for el in modality_dict["КТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "МРТ",
                "data": [el for el in modality_dict["МРТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "ММГ",
                "data": [el for el in modality_dict["ММГ"].values()],
                "barThickness": 30,
            },
            {
                "label": "Флг",
                "data": [el for el in modality_dict["Флг"].values()],
                "barThickness": 30,
            },
            {
                "label": "Флг",
                "data": [el for el in modality_dict["Флг"].values()],
                "barThickness": 30,
            },
            {
                "label": "Денс",
                "data": [el for el in modality_dict["Денс"].values()],
                "barThickness": 30,
            },
            {
                "label": "НДКТ",
                "data": [el for el in modality_dict["НДКТ"].values()],
                "barThickness": 30,
            },
            {
                "label": "ПЭТ/КТ",
                "data": [el for el in modality_dict["ПЭТ/КТ"].values()],
                "barThickness": 30,
            },
        ],
    }


@forecast_router.get("/statistics_forecast")
async def get_statistics_forecast(
    statistics_id: int,
    forecast_id: int,
    token: Annotated[str, Depends(oauth2_scheme)],
    modality: Modality = Query(None, description="Модальность"),
    # locality: Locality = Query(None, description="Локализация"),
    locality: Optional[Locality] = None,
    session: AsyncSession = Depends(get_session),
):
    """Возвращает факт-прогноз количества исследований по неделям."""
    loc = {'t1': '-', 't2': 'КУ1', 't3': 'КУ2+'}
    if locality:
        locality_value = loc.get(locality.name, None)
    else:
        locality_value = None
    # уникальные даты статистики
    query = (
        select(distinct(Statistic.date))
        .where(Statistic.id_filename == statistics_id)
        .order_by(Statistic.date.asc())
    )
    if modality:
        query = query.where(Statistic.modality == modality.value)
    if locality_value:
        query = query.where(Statistic.type == locality_value)
    query_result = await session.execute(query)
    statistics_unique_dates = query_result.scalars().all()

    # уникальные даты прогноза
    query = (
        select(distinct(Forecast.date))
        .where(Forecast.id_filename == forecast_id)
        .order_by(Forecast.date.asc())
    )
    if modality:
        query = query.where(Forecast.modality == modality.value)
    if locality_value:
        query = query.where(Forecast.type == locality_value)
    query_result = await session.execute(query)
    forecast_unique_dates = query_result.scalars().all()

    unique_dates = statistics_unique_dates + forecast_unique_dates

    # статистика
    query = (
        select(Statistic)
        .where(Statistic.id_filename == statistics_id)
        .order_by(Statistic.date.asc())
    )
    if modality:
        query = query.where(Statistic.modality == modality.value)
    if locality_value:
        query = query.where(Statistic.type == locality_value)
    query_result = await session.execute(query)
    statistics_data = query_result.scalars().all()

    # прогноз
    query = select(Forecast).where(
        Forecast.id_filename == forecast_id,
    )
    if modality:
        query = query.where(Forecast.modality == modality.value)
    if locality_value:
        query = query.where(Forecast.type == locality_value)
    query_result = await session.execute(query)
    forecast_data = query_result.scalars().all()

    count_statistics_dict = {}
    count_forecast_dict = {}

    for el in statistics_data:
        count_statistics_dict[el.date] = (
            count_statistics_dict.get(el.date, 0) + el.number
        )

    for el in forecast_data:
        count_forecast_dict[el.date] = count_forecast_dict.get(el.date, 0) + el.number

    return {
        "labels": unique_dates,
        "datasets": [
            {
                "label": "statistics",
                "data": [
                    {"x": key, "y": value}
                    for key, value in count_statistics_dict.items()
                ],
            },
            {
                "label": "forecast",
                "data": [
                    {"x": key, "y": value} for key, value in count_forecast_dict.items()
                ],
            },
        ],
    }
